package ao.ai.ml.demo.supervised.regression;

import ao.ai.ml.algo.supervised.model.algo.plain.SupervisedPlainLearner;
import ao.ai.ml.algo.supervised.model.example.ext.RegressionExample;
import ao.ai.ml.algo.supervised.model.hypothesis.SupervisedHypothesis;
import ao.ai.ml.algo.supervised.regression.linear.parametric.NormalEquationSolver;
import ao.ai.ml.demo.supervised.common.ext.RegressionProvider;
import ao.ai.ml.model.feature_set.ext.num.NumericalFeatureList;
import ao.ai.ml.model.feature_set.ext.num.SingleNumericalFeature;
import ao.ai.ml.model.feature_set.impl.FeatureScalar;
import ao.util.async.Condition;
import ao.util.math.rand.Rand;
import ao.util.misc.Uuid;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.swing.*;
import java.awt.*;
import java.awt.event.MouseEvent;
import java.awt.event.MouseListener;
import java.awt.event.MouseMotionListener;
import java.lang.reflect.InvocationTargetException;
import java.util.*;
import java.util.List;

/**
 * User: aostrovsky
 * Date: 5-Feb-2010
 * Time: 11:45:22 PM
 */
public class RegressionPanel
        extends    JPanel
        implements RegressionProvider<
                        NumericalFeatureList,
                        SingleNumericalFeature,
                        RegressionExample>,
                   MouseListener,
                   MouseMotionListener/*,
                   KeyListener*/
{
    //--------------------------------------------------------------------
    private static final Logger LOG =
            LoggerFactory.getLogger( RegressionPanel.class );

    private static final Color[] COLOURS = new Color[]{
            Color.YELLOW, Color.WHITE, Color.GREEN,
            Color.BLUE, Color.RED, Color.CYAN};

    private static final int     MAX_POINTS = 10000;


    //--------------------------------------------------------------------
    public static void main(String[] args)
            throws InvocationTargetException, InterruptedException
    {
        final RegressionPanel regression =
                new RegressionPanel();

        JFrame frame = new JFrame();
        frame.setDefaultCloseOperation( JFrame.EXIT_ON_CLOSE );
        frame.setContentPane( regression );
        frame.pack();
        frame.setSize(800, 500);
        frame.setVisible( true );

//        frame.add

        final Uuid approx = new Uuid();
        while (true)
        {
            List<XyExample> sample = new ArrayList<XyExample>();
//            for (int i = 0; i < 800; i += 10)
//            {
//                sample.add(new XyExample(
//                        i, (int) (100 * (Math.cos(
//                                Math.PI * i / 180) + 1.0)) + 10));
//            }
//            display(approx, regression, sample);

            sample.addAll( regression.nextSample() );

            final Collection<XyExample> predicted =
                    predicted(regression, sample);

            SwingUtilities.invokeAndWait(new Runnable() {
                @Override public void run() {
                    regression.display(approx, predicted);
//                    display(approx, regression, sample);
                }
            });
        }
    }

//    private static void display(
//            Uuid approx,)
//    {
//        regression.display(approx, predicted);
//    }

    private static Collection<XyExample> predicted(
            RegressionPanel                   regression,
            List<? extends RegressionExample> sample)
    {
        if (sample.size() < 2) {
            return Collections.emptyList();
        }

        SupervisedPlainLearner<
                    NumericalFeatureList,
                    SingleNumericalFeature> learner =
//                new BatchGradientDescent();
//                new StochasticGradientDescent();
                new NormalEquationSolver();
//                new LocallyWeightedRegression();

        SupervisedHypothesis<
                NumericalFeatureList,
                SingleNumericalFeature>
            hype = learner.learn( sample );

        Collection<XyExample> predicted =
                new ArrayList<XyExample>();
        for (int w = 0; w < regression.getWidth(); w += 5)
        {
            try {
                predicted.add(new XyExample(w, (int)
                       Math.round(hype.predict(new FeatureScalar(
                               w, null)).doubleValue())));
            } catch (Throwable ignored) {}
        }
        return predicted;
    }


    //--------------------------------------------------------------------
    private final Uuid                      inputId;
    private final Map     <Uuid, Color>     colours;
    private final Multimap<Uuid, XyExample> points;

    private final Condition                 available;

    private int       nextColour    = 0;
    private XyExample hoverLocation = null;


    //--------------------------------------------------------------------
    public RegressionPanel()
    {
        inputId = new Uuid();
        colours = new HashMap<Uuid, Color>();
        points  = HashMultimap.create();

        available = new Condition(false);

        setPreferredSize(new Dimension(
                200, 200));

        addMouseListener( this );
        addMouseMotionListener( this );
//        addKeyListener( this );
    }


    //--------------------------------------------------------------------
    public synchronized void clear()
    {
        points.clear();
    }


    //--------------------------------------------------------------------
    public List<XyExample> input()
    {
        return values(inputId);
    }

    public synchronized List<XyExample> values(Uuid id)
    {
        return Lists.newArrayList( points.get(id) );
    }

    @Override
    public List<XyExample> nextSample()
    {
        final List<XyExample> nextSample =
                new ArrayList<XyExample>();
        try {
            Thread t = new Thread(new Runnable() {
                @Override public void run() {
                    try {
                        available.waitForTrue();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }

                    synchronized (RegressionPanel.this) {
                        nextSample.addAll( input() );
                        available.setFalse();
                    }
                }
            });
            t.start();
            t.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        return nextSample;
    }


    //--------------------------------------------------------------------
    public synchronized void display(
            Uuid                  id,
            Collection<XyExample> idPoints)
    {
        colour( id );
        points.removeAll( id );
        points.putAll(id, idPoints);

        SwingUtilities.invokeLater(new Runnable() {
            @Override public void run() {
                repaint();
            }
        });
    }


    //--------------------------------------------------------------------
    private synchronized Color colour(Uuid id)
    {
        Color colour = colours.get( id );
        if (colour == null)
        {
            colour = nextColour();
            colours.put(id, colour);
        }
        return colour;
    }

    private Color nextColour()
    {
        if (COLOURS.length > nextColour)
        {
            return COLOURS[ nextColour++ ];
        }
        else
        {
            return new Color(
                    128 + Rand.nextInt(128),
                    128 + Rand.nextInt(128),
                    128 + Rand.nextInt(128)
            );
        }
    }


    //--------------------------------------------------------------------
    @Override
    public synchronized void paint(Graphics graphics)
    {
//        LOG.info("colours: {}", colours);

        graphics.setColor( Color.BLACK );
        graphics.fillRect(
                0, 0, getWidth(), getHeight());

        for (Uuid id : colours.keySet())
        {
//            LOG.info("points: {}", values( id ));

            graphics.setColor( colour(id) );

            for (XyExample point : values( id ))
            {
                graphics.drawRect(
                        (int) Math.round(point.input().doubleValue()),
                        (int) Math.round(point.output().doubleValue()),
                        2, 2);
            }
        }
    }


    //-------------------------------------------------------------------
    @Override public void mousePressed (MouseEvent e) {}
    @Override public void mouseReleased(MouseEvent e) {}
    @Override public void mouseEntered (MouseEvent e) {}
    @Override public void mouseExited  (MouseEvent e) {}
    @Override public void mouseMoved   (MouseEvent e) {}

    @Override public void mouseClicked (MouseEvent e) {
        if (e.getClickCount() >= 3)
        {
            clear();
        }
        
        userInput( e );
    }
    @Override public void mouseDragged(MouseEvent e) {
        userInput( e );
    }

    private void userInput(MouseEvent e)
    {
        LinkedList<XyExample> allExamples =
                new LinkedList<XyExample>();

        allExamples.addAll( input() );
        allExamples.add( new XyExample(e.getX(), e.getY()) );

        while (allExamples.size() > MAX_POINTS)
        {
            allExamples.removeFirst();
        }

        display(inputId, allExamples);
        available.setTrue();
    }

//    @O/*verride
//    public void keyTyped(KeyEvent e)
//    {
//        clear();
//    }
//
//    @Override public void keyPressed(KeyEvent e)
//    {
//        clear();
//    }
//
//    @Override public void keyReleased(KeyEvent e)
//    {
//        clear();
//    }*/
}
